"""
python fujian1_arima_predict.py
"""

import os
import json
import numpy as np
import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
from datetime import timedelta, datetime
from tqdm import tqdm

# 文件路径
input_dir = 'fujian/fujian1/preprocessed'
output_dir = 'fujian/fujian1/predict'

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 预测日期范围
start_date = datetime(2023, 5, 16)
end_date = datetime(2023, 5, 30)

# 处理每个 JSON 文件
json_files = [f for f in os.listdir(input_dir) if f.endswith('.json')]
for filename in tqdm(json_files, desc="Processing JSON files"):
    with open(os.path.join(input_dir, filename), 'r') as f:
        json_data = json.load(f)

    # 提取日期和 qty 数据
    dates = [entry['date'] for entry in json_data]
    qty_data = [entry['qty'] for entry in json_data]

    # 将 qty 数据转换为 Pandas Series
    qty_series = pd.Series(qty_data, index=pd.to_datetime(dates))

    # 确保数据是连续的
    qty_series = qty_series.asfreq('D')  # 以天为频率重采样，确保每一天都有数据

    # 拟合 ARIMA 模型
    try:
        # model = ARIMA(qty_series, order=(1, 1, 1))  # 使用较简单的模型
        model = ARIMA(qty_series, order=(1, 1, 0))  # 将 MA 阶数 q 设置为 0
        model_fit = model.fit()

        # 进行预测
        predictions = model_fit.forecast(steps=(end_date - start_date).days + 1)

        # 创建输出数据
        output_data = []
        for i in range(len(predictions)):
            date = (start_date + timedelta(days=i)).strftime('%Y-%m-%d')
            output_data.append({
                "seller_no": json_data[0]['seller_no'],
                "product_no": json_data[0]['product_no'],
                "warehouse_no": json_data[0]['warehouse_no'],
                "date": date,
                # "qty": float(predictions[i])  # 将预测值转换为 Python 原生的 float
                "qty": float(predictions.iloc[i])  # 使用 .iloc 按位置访问

            })

        # 保存预测结果
        output_filename = os.path.join(output_dir, f'predicted_{filename}')
        with open(output_filename, 'w') as f:
            json.dump(output_data, f, indent=4)

    except Exception as e:
        print(f"Error processing {filename}: {e}")

print("预测完成，结果已保存。")
